{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "import torch\n",
    "from torchvision import models\n",
    "from utils import save_net,load_net\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class CSRNet(nn.Module):\n",
    "    def __init__(self, load_weights=False):\n",
    "        super(CSRNet, self).__init__()\n",
    "        self.frontend_feat = [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512]\n",
    "        self.backend_feat  = [512, 512, 512,256,128,64]\n",
    "        self.frontend = make_layers(self.frontend_feat)\n",
    "        self.backend = make_layers(self.backend_feat,in_channels = 512,dilation = True)\n",
    "        self.output_layer = nn.Conv2d(64, 1, kernel_size=1)\n",
    "        if not load_weights:\n",
    "            mod = models.vgg16(pretrained = True)\n",
    "            self._initialize_weights()\n",
    "            for i in xrange(len(self.frontend.state_dict().items())):\n",
    "                self.frontend.state_dict().items()[i][1].data[:] = mod.state_dict().items()[i][1].data[:]\n",
    "    def forward(self,x):\n",
    "        x = self.frontend(x)\n",
    "        x = self.backend(x)\n",
    "        x = self.output_layer(x)\n",
    "        return x\n",
    "    def _initialize_weights(self):\n",
    "        for m in self.modules():\n",
    "            if isinstance(m, nn.Conv2d):\n",
    "                nn.init.normal_(m.weight, 0.01)\n",
    "                if m.bias is not None:\n",
    "                    nn.init.constant_(m.bias, 0)\n",
    "            elif isinstance(m, nn.BatchNorm2d):\n",
    "                nn.init.constant_(m.weight, 1)\n",
    "                nn.init.constant_(m.bias, 0)\n",
    "            elif isinstance(m, nn.Linear):\n",
    "                nn.init.normal_(m.weight, 0, 0.01)\n",
    "                nn.init.constant_(m.bias, 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def make_layers(cfg, in_channels = 3,batch_norm=False,dilation = False):\n",
    "    if dilation:\n",
    "        d_rate = 2\n",
    "    else:\n",
    "        d_rate = 1\n",
    "    layers = []\n",
    "    for v in cfg:\n",
    "        if v == 'M':\n",
    "            layers += [nn.MaxPool2d(kernel_size=2, stride=2)]\n",
    "        else:\n",
    "            conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=d_rate,dilation = d_rate)\n",
    "            if batch_norm:\n",
    "                layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]\n",
    "            else:\n",
    "                layers += [conv2d, nn.ReLU(inplace=True)]\n",
    "            in_channels = v\n",
    "    return nn.Sequential(*layers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "model = CSRNet()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "x = torch.rand((1,3,255,255))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 1, 31, 31])"
      ]
     },
     "execution_count": 55,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model(x).shape"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
